package vpc.ext.vstack;

import vpc.tir.*;
import vpc.tir.opt.TIRCallShape;
import vpc.tir.opt.DepthFirstOpVisitor;
import vpc.tir.expr.Operator;
import vpc.tir.TIRConst;
import vpc.core.Program;
import vpc.core.types.TypeRef;
import vpc.core.types.Type;
import vpc.core.decl.Method.Temporary;
import vpc.core.decl.Method;
import vpc.core.decl.Method.Temporary;
import vpc.core.ProgramDecl.EntryPoint;
import vpc.core.types.Type;
import vpc.hil.*;
import vpc.hil.DeviceDecl.GetRegister;
import vpc.hil.DeviceDecl.SetRegister;
import vpc.ext.vstack.util.DataHandler;
import vpc.ext.vstack.util.Debug;
import vpc.ext.vstack.util.VstackException;

import java.util.*;

public class GraphVisitor extends
		DepthFirstOpVisitor<TIRExpr, GraphVisitor.Context> {

	private static Debug.Flag debugFlag = Debug.Flag.GRAPHVISITOR;

	/**
	 * The Context class represents information that is passed from a parent
	 * node to a child node. It enables the child node to read information about
	 * its parent.
	 * 
	 * @author Philip Russell
	 * 
	 */
	public class Context {
		Node root;

		Node parent;

		public Context(Node parent) {
			this.parent = parent;
		}

		public Context(Node root, Node parent) {
			this.root = root;
			this.parent = parent;
		}
	}

	// A handle 
	final Context rootContext;
 
	final Program program;

	/**
	 * Constructor for the GraphVisitor class. Sets up the root node in the graph.
	 * @param p
	 */
	public GraphVisitor(Program p) {
		super(null);
		Method main = p.programDecl.mainEntry.method;
		Node root = new Node(main.getFullName(), DataHandler.getStackPush(main),
				p.targetDevice);
		this.rootContext = new Context(root, root);
		this.program = p;
	}

	/**
	 * Visit method for TIRExpr that is returned as the result of the getBody()
	 * method.
	 * 
	 * @param body
	 *            TIRExpr representing the body of a method.
	 * @return
	 */
	public TIRExpr visitBody(TIRExpr body, Context c) {
		return body.accept(this, c);
	}

	private void printInterrupts(Register r, Set<Interrupt> list) {
		Debug.println("REGISTER : " + r.getName(), debugFlag);
		Debug.print("( ", debugFlag);
		if (list != null) {
			for (Interrupt i : list) {
				Debug.println(i.getName() + ",", debugFlag);
			}
		}
		Debug.println(")", debugFlag);

		return;
	}

	/**
	 * In this method, we visit a known SetRegister operation in the TIR
	 * Expression tree. A SetRegister is the only operation in Virgil that can
	 * change the state of the interrupt mask. This method either statically
	 * determines the set of enabled/ disabled interrupts (if the argument is a
	 * raw, constant value) or it conservatively assumes that all interrupts are
	 * enabled by the SetRegister operation.
	 * 
	 * @param setreg
	 *            The TIRExpr object that represents writing to a register
	 * @param operands
	 *            The sub-tree of TIRExpr(s) that represent the values to be
	 *            written
	 * @return Null
	 */
	public TIRExpr visit(SetRegister setreg, TIRExpr[] operands, Context c) {

		// Get the register that is written to
		Register register = setreg.register;

		// Check to see if the target register belongs to the target device.
		if (InterruptMask.getInterrupts(register) == null) {
			// In this case, we are setting a register that is NOT
			// specified in the program's target device HIL file.
			// Most likey, this is due to an operation occuring in a library.
			// In any case, there ARE cases of SetRegister expressions that
			// do not affect a register belonging to the target device.
			// Hence, we cannot determine the set of affected interrupts.
			// TODO: is the information for these registers located anywhere
			// else
			// in the compiler? Or is it determined at link-time?
			return null;
		}
		// Else, the register is part of the target device.

		// Get a copy of the current mask.
		InterruptMask newMask = new InterruptMask(c.parent.mask);

		assert (newMask != null);

		// a SetReg operation is a UNARY operator; hence, it should only accept
		// one sub-expression.
		assert (operands.length == 1);
		// Get the sole operand
		TIRExpr operand = operands[0];

		// If the operation takes a raw, constant argument that we can
		// analyze at compile-time:
		if (operand instanceof TIRConst.Value) {
			// Given the constant value (TIRConst.Value), get the list
			// of interrupts that can can be enabled by setting the register
			Set<Interrupt> enabledInterrupts = InterruptMask.getEnabled(
					register, ((TIRConst.Value) operand).getValue());

			for (Interrupt i : InterruptMask.getInterrupts(register)) {
				if (enabledInterrupts.contains(i)) {
					newMask.enableInterrupt(i);
				} else {
					newMask.disableInterrupt(i);
				}
			}

		} else {
			// The argument to the register is a non-constant/raw value
			// Conservatively assume that each interrupt that CAN be
			// enabled by writing to the register IS enabled.
			for (Interrupt i : InterruptMask.getInterrupts(register)) {
				newMask.enableInterrupt(i);
			}
		}

		// Create a new node state that has the same name as the parent, a zero
		// weight and the newly computed mask.
		// NOTE: the weight is zero because there is nothing pushed onto the
		// stack at the point where the mask is changed, only when the jump to
		// the handler is made.
		Node newState = new Node(c.parent.name, 0, newMask);
		c.parent.addEdge(newState);
		c.parent = newState;

		// If the master bit is enabled, then iterate through the current mask.
		// For each enabled interrupt, create an edge to its handler (recurse the analysis on the handler).
		if (newMask.masterBitIsEnabled()) {
			for (Interrupt i : newMask.getInterrupts()) {
				if (newMask.isEnabled(i)) {

					Method handler = InterruptMask.getInterruptHandler(i,
							this.program);

					InterruptMask handlerMask = new InterruptMask(newMask);
					handlerMask.disableMaster();

					Node handlerNode = new Node(handler.getFullName(),
							DataHandler.getStackPush(handler), handlerMask);
					Context handlerContext = new Context(handlerNode);

					this.visitMethod(handler, handlerContext);

					c.parent.addEdge(handlerNode);
				}
			}
		}

		return null;
	}

	public void visitMethod(Method m, Context c) {
		TIRRep r = TIRUtil.getRep(m);
		this.visitBody(r.getBody(), c);
		return;
	}

	public TIRExpr visitOperator(Operator o, TIRExpr[] operands, Context c) {

		if (o instanceof SetRegister) {
			this.visit((SetRegister) o, operands, c);
		}
		return null;
	}

	/**
	 * Visit method for TIROperator type. Visit each sub-operand.
	 */
	public TIRExpr visit(TIROperator o, Context c) {

		for (TIRExpr e : o.operands)
			e.accept(this, c);

		this.visitOperator(o.operator, o.operands, c);

		return null;
	}

	/**
	 * Visit method for TIRConst.Symbol type. Represents a String, essentially.
	 * {@link vpc.tir.TIRConst.Symbol}
	 */
	public TIRExpr visit(TIRConst.Symbol s, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRConst.Value type. {@link vpc.tir.TIRConst.Value}
	 */
	public TIRExpr visit(TIRConst.Value v, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRLocal.Get type. {@link vpc.tir.TIRLocal.Get}
	 */
	public TIRExpr visit(TIRLocal.Get g, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRLocal.Set type. {@link vpc.tir.TIRLocal.Set}
	 */
	public TIRExpr visit(TIRLocal.Set l, Context c) {
		l.value.accept(this, c);
		return null;
	}

	/**
	 * @author Philip Russell
	 * @author Micah Wendell
	 * @param b
	 * @return
	 */
	private List<Method> getMethodTargets(TIRCall b) {
		TIRCallShape.CallShape cs = TIRCallShape.match(program, b);

		// The List of potential methods that COULD BE the targets
		List<Method> methodsInvoked = new ArrayList<Method>();

		// If the CallShape is direct: a direct method call only has one
		// potential target
		if (cs instanceof TIRCallShape.Direct) {

			methodsInvoked.add(((TIRCallShape.Direct) cs).method);

			// Else, if the CallShape is dynamic: then the method call could
			// target a number of different
			// methods in the program.
		} else if (cs instanceof TIRCallShape.Dynamic) {

			for (Method m : program.closure.methods) {
				TypeRef[] params = m.getArgumentTypes();
				int numArgs = b.arguments.length;

				// BEGIN: tests for equality for a dynamic method
				if (params.length != numArgs)
					continue;

				if (b.getType() != m.getType())
					continue;

				for (int i = 0; i < numArgs; i++) {
					if (params[i].getType() != b.arguments[i].getType())
						continue;
				}
				// END: test for dynamic equality

				methodsInvoked.add(m);
			}
		} else { // Virtual
			// Else, the CallShape is Virtual.
			Method toCall = ((TIRCallShape.Virtual) cs).method;
			TypeRef[] toCallParams = toCall.getArgumentTypes();
			int numArgs = toCallParams.length;

			for (Method m : program.closure.methods) {
				TypeRef[] params = m.getArgumentTypes();

				if (!m.getName().equals(toCall.getName()))
					continue;

				if (params.length != numArgs)
					continue;

				if (b.getType() != m.getType())
					continue;

				for (int i = 0; i < numArgs; i++) {
					if (params[i].getType() != toCallParams[i].getType())
						continue;
				}

				methodsInvoked.add(m);
			}
		}
		return methodsInvoked;
	}

	/**
	 * Visit method for TIRCall type. NOTE: extremely important for call graph
	 * construction.
	 * 
	 * @author Philip Russell
	 * {@link vpc.tir.TIRCall}
	 */
	public TIRExpr visit(TIRCall b, Context c) {

		// Get the list of potential target methods
		List<Method> methodsInvoked = this.getMethodTargets(b);
		List<Node> targetNodes = new ArrayList<Node>();
		assert (c.parent != null);
		targetNodes.add(c.parent);

		for (Method m : methodsInvoked) {
			Node targetNode = new Node(m.getFullName(), DataHandler
					.getStackPush(m), new InterruptMask(c.parent.mask));
			Context subContext = new Context(targetNode);

			this.visitMethod(m, subContext);

			c.parent.addEdge(targetNode);

			targetNodes.add(targetNode);
		}

		if (InterruptMask.hasDevice()) {
			if (c.parent.mask.masterBitIsEnabled()) {
				for (Node node : targetNodes) {
					for (Interrupt i : c.parent.mask.getInterrupts()) {
						if (c.parent.mask.isEnabled(i)) {
							Method handler = InterruptMask.getInterruptHandler(
									i, program);
							InterruptMask maskCopy = new InterruptMask(
									c.parent.mask);
							maskCopy.disableMaster();

							Node handlerNode = new Node(handler.getFullName(),
									DataHandler.getStackPush(handler), maskCopy);

							Context subContext = new Context(handlerNode);

							this.visitMethod(handler, subContext);

							node.addEdge(handlerNode);
						}
					}
				}
			}
		}
		for (TIRExpr arg : b.arguments) {
			arg.accept(this, c);
		}

		return null;
	}

	/**
	 * Generic visit method for TIRExpr(s) not listed above.
	 * {@link vpc.tir.TIRExpr}
	 */
	public TIRExpr visit(TIRExpr e, Context c) {

		Debug.println("WARNING: visited unknown TIRExpr, CLASS = "
				+ e.getClass().toString(), debugFlag);
		return null;
	}
}
